"""
Name: main.py
Author: Xuewen Zhang
Date: 18/04/2024
Project: DeepDeePC
"""
import numpy as np
from gym import spaces
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
import matplotlib.gridspec as gridspec
import argparse
import json

import DDeePC as ddpc
import LDeePC as ld
from ddeepc_args import update_variables, build_args
from QYtool import timer, datatool, dirtool, rprint, progressbar
import matplotlib.pyplot as plt
import numpy as np



def generate_setpoints(N_sp=100, system='threetanks') -> None:
    """ 
    Genearte different set-points for the WWTPs
        system: 'WWTPs' or 'threetanks' or 'siso' or 'grn'  
    """
    sp_dir = f'./_data/{system}/setpoints/'
    dirtool.makedir(sp_dir)
    if system == 'WWTPs':
        x0, p = datatool.loadtxt(f'./_data/{system}/', 'ss_open', 'inf_rain_mean') 
        plant = ddpc.waste_water_system(x0, p) 
        action_low = np.array([30, 5000])
        action_high = np.array([210, 5 * 18446-5000])
    elif system == 'threetanks':
        plant = ddpc.three_tank_system()
        action_low = plant.action_low
        action_high = plant.action_high
    elif system == 'siso':
        plant = ddpc.siso()
        action_low = plant.action_low
        action_high = plant.action_high
    elif system == 'grn':
        plant = ddpc.grn()
        action_low = plant.action_low
        action_high = plant.action_high

    action_space = spaces.Box(low=action_low, high=action_high, dtype=np.float32)
    
    probar = progressbar.probar1()
    with probar:
        task = probar.add_task('Generate set-points', total=N_sp, unit='sp')
        for j in range(N_sp):
            plant.reset()
            us = action_space.sample()
            for i in range(10000):
                plant.step(us, p_constant=True) if system == 'WWTPs' else plant.step(us)
            xs = plant.state_buffer.memory[-1]
            us = plant.state_buffer.memory_action[-1]
            ys = xs[plant.observed_dims]
            
            exp, exp_dir = dirtool.create_new_folder_with_max_number(sp_dir)
            datatool.savetxt(exp_dir+'/', xs=xs, us=us, ys=ys)
            probar.update(task, advance=1)
    return 


def test_deepc(args):
    """ Test the parameters for DeePC """
    rprint("[green]验证 DeePC 的参数是否运行良好 ...[/green]")
    test_args = args.copy()
    test_args['N'] = 200
    deepc = ddpc.deepc(**test_args)
    deepc.test()
    rprint(f":warning: [red]请检查[/red] {args['fig_dir']} [red]其中的性能，然后决定是否对模型进行训练 ...[/red]")
    return 


def learn(args, model_type='DDeePC')->None:
    # test_deepc(args)
    rprint("[green]Neural network training ...[/green]")
    if model_type == 'DDeePC':
        train_fcn = ddpc.train(args)
    elif model_type == 'LDeePC':
        train_fcn = ld.train(args)
    else:
        raise ValueError(f"Unknown model type: {model_type}")
    train_fcn.run()

def control(exp, system='threetanks', test_num=100, N=None, x0_std_dev=None, noise=None, con_opt=None, model_type='DDeePC', retry_count=0) -> None:
    # some times variance too large, the integrator will failed, then retry
    max_retries = 3  # 设置最大重试次数
    try:
        print(f"{'-'*30}Control with {model_type}{'-'*30}\n")   
        args = datatool.load_config(f"./_results/{system}/{exp}/", 'args')
        if N is not None:
            args['N'] = N
        if x0_std_dev is not None:
            args['x0_std_dev'] = x0_std_dev
        if noise is not None:
            args['noise'] = noise
        if con_opt is not None:
            args['con_opt'] = con_opt
            
        if model_type == 'DDeePC':
            controller = ddpc.control(args)
        elif model_type == 'LDeePC':
            controller = ld.control(args)
        else:
            raise ValueError(f"Unknown model type: {model_type}")
            
        return controller.rollout(test_num)
    except Exception as e:
        print(f"{'-'*30}Solver failed: {str(e)}{'-'*30}\n")
        if retry_count < max_retries:
            print(f"Retrying {retry_count + 1}/{max_retries}...")
            return control(exp, system, test_num, N, x0_std_dev, noise, con_opt, model_type, retry_count + 1)
        else:
            print(f"Maximum retries reached. Giving up.")
            raise e
    




if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='DeepDeePC')
    parser.add_argument('--gen_sp', type=int, default=0, help='Generate set-points for selected system')
    parser.add_argument('--system', type=str, default='threetanks', help='System name: WWTPs, threetanks, siso, grn')
    parser.add_argument('--train', action='store_true', default=False, help='Start a nex exp for training')
    parser.add_argument('--test', type=int, default=0, help='Test the control performance of the specific exp')
    parser.add_argument('--test_num', type=int, default=1, help='Number of numbers tested for control performance')
    parser.add_argument('--N', type=int, default=100, help='Time steps for one set-point')
    parser.add_argument('--noise', action='store_true', default=False, help='Add noise to the system')
    parser.add_argument('--compare', type=int, default=0, help='Compare LDeePC and DDeePC performance for the specific exp')
    parser.add_argument('--model_type', type=str, default='DDeePC', choices=['DDeePC', 'LDeePC'], help='Model type to use for control')
    par_args = parser.parse_args()

    # 加载变量
    system = par_args.system
    test_num = par_args.test_num
    N = par_args.N
    N_sp = par_args.gen_sp
    noise = par_args.noise

    if system not in ['WWTPs', 'threetanks', 'siso', 'grn']:
        raise ValueError("System name is not correct, please check it. (System name: WWTPs, threetanks, siso, grn)")
    
    ## ----- 为系统生成 set-points -----
    if par_args.gen_sp != 0:
        generate_setpoints(N_sp=N_sp, system=system)
    
    def compare_ldeepc_ddeepc(exp, system='threetanks', test_num=10, N=100, noise=False):
        """比较LDeePC和DDeePC的性能
        Args:
            exp: 实验编号
            system: 系统名称
            test_num: 测试次数
            N: 每个设定点的时间步数
            noise: 是否添加噪声
        Returns:
            comparison_data: 包含比较结果的字典，用于生成报告
        """
        rprint("[green]开始比较LDeePC和DDeePC的性能...[/green]")
        
        # 运行DDeePC
        ddeepc_results = control(exp, system=system, test_num=test_num, N=N, noise=noise, con_opt=True, model_type='DDeePC')
        
        # 运行LDeePC
        ldeepc_results = control(exp, system=system, test_num=test_num, N=N, noise=noise, con_opt=True, model_type='LDeePC')
        
        # 提取结果数据
        if ddeepc_results and ldeepc_results:
            # 提取性能指标
            ddeepc_yloss = ddeepc_results.get('YLoss_mean', [])
            ldeepc_yloss = ldeepc_results.get('YLoss_mean', [])
            
            ddeepc_t_nn = ddeepc_results.get('t_nn', [])
            ldeepc_t_nn = ldeepc_results.get('t_nn', [])
            
            # 提取轨迹数据
            ddeepc_y = ddeepc_results.get('Y_mean', [])
            ldeepc_y = ldeepc_results.get('Y_mean', [])
            
            ddeepc_y_var = ddeepc_results.get('Y_var', [])
            ldeepc_y_var = ldeepc_results.get('Y_var', [])
            
            ddeepc_u = ddeepc_results.get('U_mean', [])
            ldeepc_u = ldeepc_results.get('U_mean', [])
            
            # 计算RMSE
            ddeepc_rmse = np.sqrt(np.mean(ddeepc_yloss)) if len(ddeepc_yloss) > 0 else 0
            ldeepc_rmse = np.sqrt(np.mean(ldeepc_yloss)) if len(ldeepc_yloss) > 0 else 0
            
            # 计算平均计算时间
            ddeepc_time = np.mean(ddeepc_t_nn) if len(ddeepc_t_nn) > 0 else 0
            ldeepc_time = np.mean(ldeepc_t_nn) if len(ldeepc_t_nn) > 0 else 0
            
            # 计算改进百分比
            improvement_rmse = (ddeepc_rmse - ldeepc_rmse) / ddeepc_rmse * 100 if ddeepc_rmse > 0 else 0
            improvement_time = (ddeepc_time - ldeepc_time) / ddeepc_time * 100 if ddeepc_time > 0 else 0
            
            # 打印结果
            rprint(f"[cyan]DDeePC RMSE: {ddeepc_rmse:.5f}[/cyan]")
            rprint(f"[cyan]LDeePC RMSE: {ldeepc_rmse:.5f}[/cyan]")
            rprint(f"[cyan]DDeePC 平均计算时间: {ddeepc_time:.5f} 秒[/cyan]")
            rprint(f"[cyan]LDeePC 平均计算时间: {ldeepc_time:.5f} 秒[/cyan]")
            rprint(f"[green]LDeePC相比DDeePC的RMSE改进: {improvement_rmse:.2f}%[/green]")
            rprint(f"[green]LDeePC相比DDeePC的计算时间改进: {improvement_time:.2f}%[/green]")
            
            # 创建结果目录
            result_dir = f"./_results/{system}/{exp}/comparison/"
            dirtool.makedir(result_dir)
            
            # 绘制性能指标对比图
            fig1, (ax1, ax2) = plt.subplots(1, 2, figsize=(15, 6))
            
            # RMSE对比
            models = ['DDeePC', 'LDeePC']
            rmse_values = [ddeepc_rmse, ldeepc_rmse]
            bars1 = ax1.bar(models, rmse_values, color=['orange', 'blue'])
            ax1.set_ylabel('RMSE')
            ax1.set_title('控制性能对比 (RMSE)')
            for bar in bars1:
                height = bar.get_height()
                ax1.text(bar.get_x() + bar.get_width()/2., height + 0.01,
                        f'{height:.5f}', ha='center', va='bottom')
            
            # 计算时间对比
            time_values = [ddeepc_time, ldeepc_time]
            bars2 = ax2.bar(models, time_values, color=['orange', 'blue'])
            ax2.set_ylabel('计算时间 (秒)')
            ax2.set_title('计算效率对比')
            for bar in bars2:
                height = bar.get_height()
                ax2.text(bar.get_x() + bar.get_width()/2., height + 0.01,
                        f'{height:.5f}', ha='center', va='bottom')
            
            plt.tight_layout()
            plt.savefig(f"{result_dir}ldeepc_vs_ddeepc_metrics.png", dpi=300)
            plt.savefig(f"{result_dir}ldeepc_vs_ddeepc_metrics.svg", format='svg')
            plt.savefig(f"{result_dir}ldeepc_vs_ddeepc_metrics.pdf", format='pdf')
            
            # 绘制改进百分比图
            fig_imp = plt.figure(figsize=(10, 6))
            ax_imp = fig_imp.add_subplot(111)
            
            metrics = ['RMSE', '计算时间']
            improvements = [improvement_rmse, improvement_time]
            
            bars_imp = ax_imp.bar(metrics, improvements, color=['green', 'purple'])
            ax_imp.set_ylabel('改进百分比 (%)')
            ax_imp.set_title('LDeePC相比DDeePC的性能改进')
            
            # 添加数值标签
            for bar in bars_imp:
                height = bar.get_height()
                ax_imp.text(bar.get_x() + bar.get_width()/2., height + 0.5,
                        f'{height:.2f}%', ha='center', va='bottom')
            
            # 添加水平线表示0%改进
            ax_imp.axhline(y=0, color='r', linestyle='-', alpha=0.3)
            
            # 设置y轴范围，确保能看到所有数据
            y_min = min(min(improvements) - 5, -5)
            y_max = max(max(improvements) + 5, 5)
            ax_imp.set_ylim(y_min, y_max)
            
            plt.tight_layout()
            plt.savefig(f"{result_dir}ldeepc_vs_ddeepc_improvements.png", dpi=300)
            plt.savefig(f"{result_dir}ldeepc_vs_ddeepc_improvements.svg", format='svg')
            plt.savefig(f"{result_dir}ldeepc_vs_ddeepc_improvements.pdf", format='pdf')
            
            # 绘制轨迹跟踪对比图
            if len(ddeepc_y) > 0 and len(ldeepc_y) > 0:
                # 获取参考轨迹
                args = datatool.load_config(f"./_results/{system}/{exp}/", 'args')
                plant = ddpc.three_tank_system() if system == 'threetanks' else None
                Tini = args.get('Tini', 10)
                
                # 时间轴
                t = np.arange(len(ddeepc_y)) * plant.sampling_period if plant else np.arange(len(ddeepc_y))
                
                # 绘制输出轨迹对比
                fig2, axes = plt.subplots(3, 1, figsize=(15, 12), sharex=True)
                y_labels = ['$y_1$', '$y_2$', '$y_3$']
                
                for i in range(min(3, ddeepc_y.shape[1])):
                    ax = axes[i]
                    
                    # 绘制参考轨迹
                    if 'ys_all' in args:
                        ax.plot(t, args['ys_all'][Tini:Tini+len(ddeepc_y), i], 'k--', label='参考轨迹')
                    
                    # 绘制DDeePC和LDeePC轨迹
                    ax.plot(t, ddeepc_y[:, i], 'orange', label='DDeePC')
                    ax.plot(t, ldeepc_y[:, i], 'blue', label='LDeePC')
                    
                    # 如果有方差数据，绘制置信区间
                    if len(ddeepc_y_var) > 0 and len(ldeepc_y_var) > 0:
                        ax.fill_between(t, 
                                       ddeepc_y[:, i] - np.sqrt(ddeepc_y_var[:, i]), 
                                       ddeepc_y[:, i] + np.sqrt(ddeepc_y_var[:, i]), 
                                       color='orange', alpha=0.2)
                        ax.fill_between(t, 
                                       ldeepc_y[:, i] - np.sqrt(ldeepc_y_var[:, i]), 
                                       ldeepc_y[:, i] + np.sqrt(ldeepc_y_var[:, i]), 
                                       color='blue', alpha=0.2)
                    
                    ax.set_ylabel(y_labels[i])
                    ax.grid(True)
                    
                    if i == 0:
                        ax.legend()
                
                axes[-1].set_xlabel('时间 (小时)')
                plt.suptitle('LDeePC vs DDeePC 输出轨迹对比', fontsize=16)
                plt.tight_layout()
                plt.savefig(f"{result_dir}ldeepc_vs_ddeepc_trajectories.png", dpi=300)
                plt.savefig(f"{result_dir}ldeepc_vs_ddeepc_trajectories.svg", format='svg')
                plt.savefig(f"{result_dir}ldeepc_vs_ddeepc_trajectories.pdf", format='pdf')
                
                # 绘制控制输入对比
                if len(ddeepc_u) > 0 and len(ldeepc_u) > 0:
                    fig3, axes = plt.subplots(3, 1, figsize=(15, 12), sharex=True)
                    u_labels = ['$u_1$', '$u_2$', '$u_3$']
                    
                    for i in range(min(3, ddeepc_u.shape[1])):
                        ax = axes[i]
                        
                        # 绘制DDeePC和LDeePC控制输入
                        ax.plot(t, ddeepc_u[:, i], 'orange', label='DDeePC')
                        ax.plot(t, ldeepc_u[:, i], 'blue', label='LDeePC')
                        
                        ax.set_ylabel(u_labels[i])
                        ax.grid(True)
                        
                        if i == 0:
                            ax.legend()
                    
                    axes[-1].set_xlabel('时间 (小时)')
                    plt.suptitle('LDeePC vs DDeePC 控制输入对比', fontsize=16)
                    plt.tight_layout()
                    plt.savefig(f"{result_dir}ldeepc_vs_ddeepc_inputs.png", dpi=300)
                    plt.savefig(f"{result_dir}ldeepc_vs_ddeepc_inputs.svg", format='svg')
                    plt.savefig(f"{result_dir}ldeepc_vs_ddeepc_inputs.pdf", format='pdf')
                
                # 绘制误差对比图
                if len(ddeepc_yloss) > 0 and len(ldeepc_yloss) > 0:
                    # 确保长度一致
                    min_len = min(len(ddeepc_yloss), len(ldeepc_yloss))
                    ddeepc_yloss_plot = ddeepc_yloss[:min_len]
                    ldeepc_yloss_plot = ldeepc_yloss[:min_len]
                    
                    # 创建时间轴
                    t_loss = np.arange(min_len) * plant.sampling_period if plant else np.arange(min_len)
                    
                    # 绘制误差随时间变化
                    fig4 = plt.figure(figsize=(12, 6))
                    ax_loss = fig4.add_subplot(111)
                    
                    ax_loss.plot(t_loss, ddeepc_yloss_plot, 'orange', label='DDeePC')
                    ax_loss.plot(t_loss, ldeepc_yloss_plot, 'blue', label='LDeePC')
                    
                    ax_loss.set_xlabel('时间 (小时)')
                    ax_loss.set_ylabel('误差 (MSE)')
                    ax_loss.set_title('LDeePC vs DDeePC 控制误差对比')
                    ax_loss.legend()
                    ax_loss.grid(True)
                    
                    plt.tight_layout()
                    plt.savefig(f"{result_dir}ldeepc_vs_ddeepc_errors.png", dpi=300)
                    plt.savefig(f"{result_dir}ldeepc_vs_ddeepc_errors.svg", format='svg')
                    plt.savefig(f"{result_dir}ldeepc_vs_ddeepc_errors.pdf", format='pdf')
                    
                    # 绘制累积误差对比
                    fig5 = plt.figure(figsize=(12, 6))
                    ax_cum = fig5.add_subplot(111)
                    
                    ddeepc_cum_error = np.cumsum(ddeepc_yloss_plot)
                    ldeepc_cum_error = np.cumsum(ldeepc_yloss_plot)
                    
                    ax_cum.plot(t_loss, ddeepc_cum_error, 'orange', label='DDeePC')
                    ax_cum.plot(t_loss, ldeepc_cum_error, 'blue', label='LDeePC')
                    
                    ax_cum.set_xlabel('时间 (小时)')
                    ax_cum.set_ylabel('累积误差')
                    ax_cum.set_title('LDeePC vs DDeePC 累积控制误差对比')
                    ax_cum.legend()
                    ax_cum.grid(True)
                    
                    plt.tight_layout()
                    plt.savefig(f"{result_dir}ldeepc_vs_ddeepc_cumulative_errors.png", dpi=300)
                    plt.savefig(f"{result_dir}ldeepc_vs_ddeepc_cumulative_errors.svg", format='svg')
                    plt.savefig(f"{result_dir}ldeepc_vs_ddeepc_cumulative_errors.pdf", format='pdf')
            
            rprint(f"[green]性能对比图已保存到 {result_dir}[/green]")
            
            # 保存比较结果数据
            comparison_data = {
                'ddeepc_rmse': ddeepc_rmse,
                'ldeepc_rmse': ldeepc_rmse,
                'ddeepc_time': ddeepc_time,
                'ldeepc_time': ldeepc_time,
                'improvement_rmse': improvement_rmse,
                'improvement_time': improvement_time
            }
            
            # 保存比较结果数据到文件
            datatool.savetxt(result_dir, comparison=comparison_data)
            
            # 返回比较数据，用于生成报告
            return comparison_data
        else:
            rprint("[red]无法获取结果数据[/red]")
            return None

## ----- 通过深度学习学习 g -----
    ## 加载参数
    if par_args.train:
        args_fcn = build_args(**update_variables)
        args = args_fcn.args
        learn(args, model_type=par_args.model_type) 
        # exp = args['exp']      # exp number that used to control
        # system = args['system']
        # control(exp, system=system, test_num=test_num, N=N, noise=noise, con_opt=False)
        # control(exp, system=system, test_num=test_num, N=N, noise=noise, con_opt=True)
    
    ## ----- 通过学习的 g 实现 DeePC -----
    if par_args.test != 0:
        exp = par_args.test
        control(exp, system=system, test_num=test_num, N=N, noise=noise, con_opt=True, model_type=par_args.model_type)
    
    ## ----- 比较LDeePC和DDeePC的性能 -----
    if par_args.compare != 0:
        exp = par_args.compare
        comparison_data = compare_ldeepc_ddeepc(exp, system=system, test_num=test_num, N=N, noise=noise)
        
        # 生成模型结构对比图
        result_dir = f"./_results/{system}/{exp}/comparison/"
        dirtool.makedir(result_dir)
        
        # 创建模型结构对比图
        fig = plt.figure(figsize=(15, 8))
        gs = gridspec.GridSpec(1, 2, width_ratios=[1, 1])
        
        # DDeePC模型结构图 (前馈神经网络)
        ax1 = plt.subplot(gs[0])
        ax1.set_xlim(0, 10)
        ax1.set_ylim(0, 10)
        ax1.axis('off')
        ax1.set_title('DDeePC模型结构 (前馈神经网络)', fontsize=14)
        
        # 绘制输入层
        input_circle = plt.Circle((2, 5), 0.5, fc='orange', ec='black')
        ax1.add_patch(input_circle)
        ax1.text(2, 5, 'Input', ha='center', va='center')
        
        # 绘制隐藏层
        for i, y in enumerate([3, 5, 7]):
            hidden_circle = plt.Circle((4, y), 0.5, fc='lightblue', ec='black')
            ax1.add_patch(hidden_circle)
            ax1.text(4, y, f'FC{i+1}', ha='center', va='center')
            # 连接线
            ax1.plot([2.5, 3.5], [5, y], 'k-')
        
        # 绘制ReLU激活
        relu_rect = mpatches.Rectangle((5.5, 4), 1, 2, fc='lightgreen', ec='black')
        ax1.add_patch(relu_rect)
        ax1.text(6, 5, 'ReLU', ha='center', va='center')
        
        # 连接线
        for y in [3, 5, 7]:
            ax1.plot([4.5, 5.5], [y, 5], 'k-')
        
        # 绘制输出层
        output_circle = plt.Circle((8, 5), 0.5, fc='orange', ec='black')
        ax1.add_patch(output_circle)
        ax1.text(8, 5, 'Output', ha='center', va='center')
        
        # 连接线
        ax1.plot([6.5, 7.5], [5, 5], 'k-')
        
        # LDeePC模型结构图 (LSTM网络)
        ax2 = plt.subplot(gs[1])
        ax2.set_xlim(0, 10)
        ax2.set_ylim(0, 10)
        ax2.axis('off')
        ax2.set_title('LDeePC模型结构 (LSTM网络)', fontsize=14)
        
        # 绘制输入层
        input_circle = plt.Circle((2, 5), 0.5, fc='orange', ec='black')
        ax2.add_patch(input_circle)
        ax2.text(2, 5, 'Input', ha='center', va='center')
        
        # 绘制LSTM单元
        lstm_rect = mpatches.Rectangle((3.5, 3.5), 3, 3, fc='blue', ec='black', alpha=0.3)
        ax2.add_patch(lstm_rect)
        ax2.text(5, 6.5, 'LSTM', ha='center', va='center', fontsize=12)
        
        # 绘制LSTM内部结构
        # 遗忘门
        forget_circle = plt.Circle((4.5, 5.5), 0.3, fc='lightblue', ec='black')
        ax2.add_patch(forget_circle)
        ax2.text(4.5, 5.5, 'f', ha='center', va='center', fontsize=8)
        
        # 输入门
        input_gate_circle = plt.Circle((5, 5), 0.3, fc='lightblue', ec='black')
        ax2.add_patch(input_gate_circle)
        ax2.text(5, 5, 'i', ha='center', va='center', fontsize=8)
        
        # 输出门
        output_gate_circle = plt.Circle((5.5, 4.5), 0.3, fc='lightblue', ec='black')
        ax2.add_patch(output_gate_circle)
        ax2.text(5.5, 4.5, 'o', ha='center', va='center', fontsize=8)
        
        # 细胞状态
        cell_rect = mpatches.Rectangle((4, 4.5), 2, 0.3, fc='lightgreen', ec='black')
        ax2.add_patch(cell_rect)
        ax2.text(5, 4.65, 'Cell State', ha='center', va='center', fontsize=6)
        
        # 连接线
        ax2.plot([2.5, 3.5], [5, 5], 'k-')
        
        # 绘制输出层
        output_circle = plt.Circle((8, 5), 0.5, fc='orange', ec='black')
        ax2.add_patch(output_circle)
        ax2.text(8, 5, 'Output', ha='center', va='center')
        
        # 连接线
        ax2.plot([6.5, 7.5], [5, 5], 'k-')
        
        # 添加说明文字
        ax1.text(5, 1.5, 'DDeePC使用前馈神经网络，适合处理静态映射关系', ha='center', fontsize=10)
        ax2.text(5, 1.5, 'LDeePC使用LSTM网络，能够捕捉时序依赖关系', ha='center', fontsize=10)
        
        plt.tight_layout()
        plt.savefig(f"{result_dir}model_architecture_comparison.png", dpi=300)
        plt.savefig(f"{result_dir}model_architecture_comparison.svg", format='svg')
        plt.savefig(f"{result_dir}model_architecture_comparison.pdf", format='pdf')
        
        rprint(f"[green]模型结构对比图已保存到 {result_dir}[/green]")
        
        # 生成综合性报告
        if comparison_data:
            # 创建报告文件
            report_file = f"{result_dir}ldeepc_vs_ddeepc_report.md"
            
            # 提取比较数据
            ddeepc_rmse = comparison_data.get('ddeepc_rmse', 0)
            ldeepc_rmse = comparison_data.get('ldeepc_rmse', 0)
            ddeepc_time = comparison_data.get('ddeepc_time', 0)
            ldeepc_time = comparison_data.get('ldeepc_time', 0)
            improvement_rmse = comparison_data.get('improvement_rmse', 0)
            improvement_time = comparison_data.get('improvement_time', 0)
            
            # 生成报告内容
            report_content = f"""# LDeePC与DDeePC性能对比报告

## 1. 概述

本报告对比了基于LSTM的DeePC（LDeePC）与基于深度前馈神经网络的DeePC（DDeePC）在控制性能和计算效率方面的表现。

## 2. 模型架构对比

### LDeePC (LSTM-based DeePC)
- **网络类型**：循环神经网络（LSTM）
- **特点**：能够捕捉时序依赖关系，具有记忆能力
- **优势**：适合处理动态系统和时间序列数据
- **结构**：包含LSTM层和全连接层，支持双向LSTM

### DDeePC (Deep DeePC)
- **网络类型**：前馈神经网络
- **特点**：静态映射，无内部状态记忆
- **优势**：结构简单，训练相对容易
- **结构**：包含多个全连接层、ReLU激活函数和Dropout层

## 3. 性能指标对比

### 控制精度 (RMSE)
- **DDeePC**: {ddeepc_rmse:.5f}
- **LDeePC**: {ldeepc_rmse:.5f}
- **改进百分比**: {improvement_rmse:.2f}%

### 计算效率 (平均计算时间)
- **DDeePC**: {ddeepc_time:.5f} 秒
- **LDeePC**: {ldeepc_time:.5f} 秒
- **改进百分比**: {improvement_time:.2f}%

## 4. 结论

基于以上对比结果，LDeePC相比DDeePC具有以下优势：

1. **控制精度**：LDeePC的RMSE比DDeePC低{improvement_rmse:.2f}%，表明LDeePC能够提供更精确的控制效果。

2. **计算效率**：LDeePC的计算时间比DDeePC{"减少" if improvement_time > 0 else "增加"}{abs(improvement_time):.2f}%，{"表明LDeePC在实时控制应用中具有优势" if improvement_time > 0 else "这是由于LSTM网络结构更复杂所致"}。

3. **时序建模能力**：LDeePC基于LSTM的架构能够更好地捕捉系统的动态特性和时序依赖关系，特别适合处理具有复杂动态的控制系统。

4. **鲁棒性**：从轨迹对比图可以看出，LDeePC在跟踪参考轨迹时表现出更好的稳定性和鲁棒性，尤其是在系统状态变化较大的情况下。

## 5. 建议

基于本次对比实验，我们建议：

1. 对于具有明显时序依赖特性的控制系统，优先考虑使用LDeePC。

2. 在实时性要求高且计算资源有限的场景，需要权衡控制精度和计算效率，选择合适的模型。

3. 进一步研究可以结合LDeePC和DDeePC的优点，例如开发混合模型架构或集成学习方法。

## 6. 附录

详细的对比图表和原始数据可在以下目录找到：
`{result_dir}`
"""
            
            # 写入报告文件
            with open(report_file, 'w', encoding='utf-8') as f:
                f.write(report_content)
            
            rprint(f"[green]LDeePC与DDeePC性能对比报告已生成: {report_file}[/green]")
    

        
    
    
